package com.luo.ai.assistant.chat.memory;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.luo.ai.assistant.chat.ChatConsts;
import com.luo.ai.assistant.conv.model.Chat;
import com.luo.ai.assistant.util.JsonUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.dao.DataAccessException;
import org.springframework.data.redis.core.*;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.util.CollectionUtils;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * Redis聊天记录实现，使用2级缓存策略：
 * <ol>
 *     <li>
 *         本地缓存：Caffeine缓存会话ID对应的消息列表，最大缓存会话数量为{@link MyRedisMemoryConfig#getMaxConvCountLocalCache()}
 *     </li>
 *     <li>
 *         Redis缓存：Redis zset(member=chat, score=timestampMills)，消息过期时间为{@link MyRedisMemoryConfig#getExpiredTimestampMills()}
 *     </li>
 *     <li>
 *         增量补充本地缓存：每次查询会话消息列表时，增量补充本地缓存（仅补充本地缓存与Redis zset间缺少的消息列表），清理本地会话消息缓存（过期消息、多余息数）
 *     </li>
 *     <li>
 *         定时清理：清理Redis缓存（Lua脚本）、清理本地缓存
 *     </li>
 * </ol>
 *
 * @author luohq
 * @date 2025-02-02 10:28
 */
public class MyRedisChatMemory_V2 implements ChatMemory {

    private static final Logger log = LoggerFactory.getLogger(MyRedisChatMemory_V2.class);

    /**
     * RedisTemplate
     */
    private RedisTemplate<String, Object> redisTemplate;

    /**
     * 会话ID到本地缓存的映射
     */
    private Cache<String, ChatsLocalCache> convIdToMessagesCache;

    /**
     * 会话ID到最大消息数的映射
     */
    private final Map<String, Integer> convIdToMaxMsgCountMap = new ConcurrentHashMap<>();

    /**
     * Redis ChatMemory配置
     */
    private MyRedisMemoryConfig myRedisMemoryConfig;

    /**
     * 默认获取最近20条聊天记录
     */
    private int lastN = ChatConsts.DEFAULT_CHAT_MEMORY_SIZE;

    public MyRedisChatMemory_V2(RedisTemplate<String, Object> redisTemplate, MyRedisMemoryConfig myRedisMemoryConfig) {
        this.redisTemplate = redisTemplate;
        this.myRedisMemoryConfig = myRedisMemoryConfig;
        //初始化本地会话缓存
        this.convIdToMessagesCache = Caffeine.newBuilder()
                //设置最大缓存的会话数量
                .maximumSize(myRedisMemoryConfig.getMaxConvCountLocalCache())
                .build();
    }

    public void setLastN(int lastN) {
        this.lastN = lastN;
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        //如果消息为空，直接返回
        if (CollectionUtils.isEmpty(messages)) {
            return;
        }

        //使用Pipeline，提升执行效率
        this.redisTemplate.executePipelined(new SessionCallback<Object>() {
            @Override
            public <K, V> Object execute(RedisOperations<K, V> redisOperations) throws DataAccessException {
                RedisOperations<String, Object> myRedisOperations = (RedisOperations<String, Object>) redisOperations;
                //1.记录出现过的convId
                myRedisOperations.opsForSet().add(redisConvIdSetKey(), conversationId);
                //2，记录会话中的消息
                for (int i = 0; i < messages.size(); i++) {
                    //待保存消息
                    Message message = messages.get(i);
                    //转换成chat - 支持序列化
                    Chat chat = ChatMessageConvertor.toChat(conversationId, message);
                    //避免时间戳相同，手动累加索引，使其呈递增趋势
                    Long curTimestamp = System.currentTimeMillis() + i;
                    //【相较于V1版本额外添加】取当前时间戳为chat ID（递增）
                    chat.setId(curTimestamp);
                    log.debug("[ChatMemory] Redis - add message[{}] at {}: {}", i, curTimestamp, JsonUtils.toJson(message));
                    //保存消息
                    myRedisOperations.opsForZSet().add(redisConvMessageKey(conversationId), chat, curTimestamp);
                }
                return null;
            }
        });
    }

    @Override
    public List<Message> get(String conversationId) {
        //记录会话ID到最大消息数的映射
        this.convIdToMaxMsgCountMap.put(conversationId, lastN);

        //当前会话消息的本地缓存
        ChatsLocalCache curConvChatsLocalCache = this.convIdToMessagesCache.get(conversationId, key -> new ChatsLocalCache());
        //本地缓存的最大chat ID（最后一条消息的缓存时间）
        Long localChatMaxTimestamp = curConvChatsLocalCache.getLastChatId();
        //1.增量补充本地缓存
        List<Chat> deltaChatList = this.redisTemplate.opsForZSet()
                .rangeByScore(this.redisConvMessageKey(conversationId), localChatMaxTimestamp + 1, System.currentTimeMillis())
                .stream()
                .map(Chat.class::cast).collect(Collectors.toList());
        curConvChatsLocalCache.add(deltaChatList);

        //2.清理本地缓存多余、过期的消息
        curConvChatsLocalCache.ltrim(lastN);
        curConvChatsLocalCache.clearExpiredChats(this.myRedisMemoryConfig.getExpiredTimestampMills());


        //如果lastN小于等于0，返回空列表
        if (0 >= lastN) {
            log.debug("[ChatMemory] Redis - get message history: lastN is {} and Return Empty Chat List!", lastN);
            return List.of();
        }

        //3.转换chat为message列表
        log.debug("[ChatMemory] Redis - get message history - message count: {}, message list: {}",
                curConvChatsLocalCache.getChatCount(), JsonUtils.toJson(curConvChatsLocalCache.getChatList()));
        return curConvChatsLocalCache.getMessageList();
    }

    @Override
    public void clear(String conversationId) {
        log.debug("[ChatMemory] Redis - clear message of conversation: {}", conversationId);
        //清理Redis中的消息
        this.redisTemplate.delete(this.redisConvMessageKey(conversationId));
        //清理本地缓存
        this.convIdToMessagesCache.invalidate(conversationId);
        this.convIdToMaxMsgCountMap.remove(conversationId);
    }

    /**
     * 会话对应的消息集合KEY（用于记录会话对应的消息集合）
     *
     * @param conversationId 会话ID
     * @return 会话对应的消息集合KEY
     */
    private String redisConvMessageKey(String conversationId) {
        return this.myRedisMemoryConfig.getKeyPrefix() + conversationId;
    }

    /**
     * 会话ID集合Key（用于记录所有的会话ID）
     *
     * @return 会话ID集合Key
     */
    private String redisConvIdSetKey() {
        return this.myRedisMemoryConfig.getKeyPrefix() + "convIds";
    }

    /**
     * 清理过期消息的Redis Lua脚本
     */
    private String CLEAR_TASK_LUA_SCRIPT = """
            local redisConvIdSetKey = KEYS[1]
            local redisConvMessageKey = KEYS[2]
            local convId = ARGV[1]
            local expireTimeMills = tonumber(ARGV[2])
            local curConvMaxMsgCount = tonumber(ARGV[3])
            
            -- 移除已过期的消息
            if expireTimeMills > 0 then
                redis.call("ZREMRANGEBYSCORE", redisConvMessageKey, 0, expireTimeMills)
            end
            
            -- 仅保留最大数量的消息，移除多余的消息
            if curConvMaxMsgCount > 0 then
                redis.call("ZREMRANGEBYRANK", redisConvMessageKey, 0, -(curConvMaxMsgCount + 1))
            end
            
            -- 清理无用的会话记录
            local remainConvMessageCount = redis.call("ZCARD", redisConvMessageKey)
            if remainConvMessageCount <= 0 then
                redis.call("SREM", redisConvIdSetKey, convId)
                redis.call("DEL", redisConvMessageKey)
            end
            
            return remainConvMessageCount
            """;

    @Scheduled(cron = "${spring.ai.chat.memory.redis.clear-task-cron:0/20 * * * * ?}")
    void clearExpireMessages() {
        log.debug("<== [ChatMemory] Redis - clear expire conv messages ==>");
        //scan选项 - 设置单次扫描返回数量（预估值）
        ScanOptions options = ScanOptions.scanOptions()
                .count(this.myRedisMemoryConfig.getMaxConvCountPerClearScan())
                .build();

        //Redis Lua脚本
        DefaultRedisScript<Long> clearTaskRedisScript = new DefaultRedisScript<>();
        clearTaskRedisScript.setScriptText(CLEAR_TASK_LUA_SCRIPT);
        clearTaskRedisScript.setResultType(Long.class);

        //会话ID集合Key
        String redisConvIdSetKey = this.redisConvIdSetKey();
        //1.扫描集合中存在的会话ID
        try (Cursor<Object> cursor = redisTemplate.opsForSet().scan(redisConvIdSetKey, options)) {
            while (cursor.hasNext()) {
                //待清理的会话ID
                String convId = String.valueOf(cursor.next());
                log.debug("<===== [ChatMemory] Redis - clear expire messages of conv: {} =====>", convId);
                //会话对应的消息集合KEY
                String redisConvMessageKey = this.redisConvMessageKey(convId);
                //过期时间点（该时间点之前的消息均已过期）
                Long expireTimeMills = this.myRedisMemoryConfig.getExpiredTimestampMills();
                //当前会话的最大历史消息数
                Integer curConvMaxMsgCount = this.convIdToMaxMsgCountMap.getOrDefault(convId, -1);
                //执行Lua脚本原子清理过期消息（避免分布式操作导致误删不为空的消息）
                Long remainConvMessageCount = this.redisTemplate.execute(clearTaskRedisScript,
                        List.of(redisConvIdSetKey, redisConvMessageKey),
                        convId, expireTimeMills, curConvMaxMsgCount);
                //彻底移除本地无用的记录
                if (0 >= remainConvMessageCount) {
                    this.convIdToMaxMsgCountMap.remove(convId);
                    this.convIdToMessagesCache.invalidate(convId);
                } else {
                    //清理本地缓存（仅当本地存在时才清理）
                    Optional.ofNullable(this.convIdToMessagesCache.getIfPresent(convId))
                            .ifPresent(chatsLocalCache -> {
                                chatsLocalCache.ltrim(curConvMaxMsgCount);
                                chatsLocalCache.clearExpiredChats(expireTimeMills);
                            });
                }
            }
        } catch (Exception e) {
            log.warn("[ChatMemory] Redis - clear expire messages exception!", e);
        }
    }
}
